import argparse
import typing
import torch
import safetensors as st
from safetensors.torch import save_file
from tqdm import tqdm


def decomp(m: torch.Tensor, r: int, return_loss=False, loss_mode: typing.Literal['L1', 'L2'] = 'L1'):
    h, w = m.shape[:2]
    u, s, v = torch.linalg.svd(m.to(torch.float32), full_matrices=False)
    u = u[:, :r]
    s = torch.diag(s[:r])
    v = v[:r, :]

    b = (u @ s).to(m.dtype)
    a = v.to(m.dtype)
    loss = torch.nan

    if return_loss:
        loss = torch.abs(m - (b @ a)).sum() / (h * w) if loss_mode == 'L1' else torch.sum((m - (b @ a)) ** 2) ** 0.5 / (h * w)

    return b, a, loss


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', '-i', type=str, required=True, help='lora file name')
    parser.add_argument('--output', '-o', type=str, required=True, help='output file name')
    parser.add_argument('--rank', '-r', type=int, required=True, help='new rank level')

    args = parser.parse_args()

    new_state_dict = {}
    with st.safe_open(args.input, framework='pt') as spt:
        keys = spt.keys()
        layers = [*{k.replace('.lora_A.weight', '').replace('.lora_B.weight', '') for k in keys}]
        total_loss = 0.0
        l_bar = tqdm(layers)
        for layer in l_bar:
            a_n, b_n = layer + '.lora_A.weight', layer + '.lora_B.weight'
            assert a_n in keys, ValueError(f"{a_n} not in old state dict")
            assert b_n in keys, ValueError(f"{b_n} not in old state dict")

            a_t, b_t = spt.get_tensor(a_n), spt.get_tensor(b_n)
            assert b_t.shape == a_t.T.shape, ValueError(f"lora_B shape: {b_t.shape} != lora_A.T shape: {a_t.T.shape}")
            d_W = b_t @ a_t
            d_B, d_A, d_L = decomp(d_W, args.rank, return_loss=True)
            total_loss += d_L.item()
            new_state_dict[a_n] = d_A.contiguous()
            new_state_dict[b_n] = d_B.contiguous()
            l_bar.set_postfix({
                'layer': layer,
                'B': b_t.shape,
                'A': a_t.shape,
                'new_A': new_state_dict[a_n].shape,
                'new_B': new_state_dict[b_n].shape,
                'decompose loss': d_L.item(),
            })

        print(f"total loss: {total_loss:.4f}, avg loss: {total_loss / len(layers):.4f}")

    save_file(new_state_dict, args.output)


if __name__ == '__main__':
    main()
